﻿#if DEBUG
//#define HMM_DEBUGGING
#endif

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.IO.Compression;
using System.Linq;

using alib.Debugging;
using alib.Dictionary;
using alib.Enumerable;
using alib.Hashing;
using alib.IO;

namespace alib.Hmm
{
	using Math = System.Math;
	using String = System.String;

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>
	/// 
	/// </summary>
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	public class TrigramHmm
	{
#if HMM_DEBUGGING
		public static TrigramHmm _this;
#endif
		public TrigramHmm(Dictionary<Emission, Double> prob_ep, ProbabilityTree[] p_tree)
		{
#if HMM_DEBUGGING
			_this = this;
#endif
			this.prob_ep = prob_ep;
			this.p_tree = p_tree;
		}
		protected TrigramHmm(_BinaryReader br)
			: this(read_1(br), read_2(br))
		{
		}
		static Dictionary<Emission, Double> read_1(_BinaryReader br)
		{
			int c = br.ReadInt32();
			var prob_ep = new Dictionary<Emission, Double>(c);
			for (int i = 0; i < c; i++)
				prob_ep.Add(new Emission(br.Read7BitEncodedInt(), br.Read7BitEncodedInt()), br.ReadDouble());
			return prob_ep;
		}
		static ProbabilityTree[] read_2(_BinaryReader br)
		{
			var pt = new ProbabilityTree();
			pt.Read(br);
			return pt.probs;
		}

		[DebuggerDisplay("{tag,nq} {P}")]
		public struct TagProb
		{
			public Double P;
			public String tag;
		}

		struct TrellisData
		{
			public static readonly TrellisData StartNode = new TrellisData { prob = 1.0 };
			public Double prob;
			public Bitag path;
#if HMM_DEBUGGING
			public int tag;
			public override String ToString()
			{
				StringIndex ut = ((StringHmm)_this).unitags;
				return String.Format("{0,6} {1} {2,4}", prob, path.ToString(), ut[tag]);
			}
#endif
		};
		sealed class TrellisColumn : SlotDictionary<Bitag, TrellisData> { }

		protected readonly Dictionary<Emission, Double> prob_ep;
		protected readonly ProbabilityTree[] p_tree;

#if HMM_TUNING
		public static Double pval = 1e-11;
		public static Double bitag_pref = 1.0;
		public static Double tritag_pref = 4.0;
#else
		const Double pval = 1e-11;
		const Double bitag_pref = 1.0;
		const Double tritag_pref = 4.0;
#endif
#if CLOSED_CLASS_TAGS
		static String[] closed_class = { "IN", "TO", "MD", "EX", "CC", "DT", "PRP", "PRP$", "POS", "$", "#" };

		public bool[] closed_class_tags;

		void setup_closed_classes()
		{
			var ut = ((StringHmm)this).unitags;
			closed_class_tags = new bool[ut.Count];
			for (int i = 0; i < ut.Count; i++)
				closed_class_tags[i] = ut[i].Length <= 1 || System.Array.IndexOf<String>(closed_class, ut[i]) != -1;
		}
#endif

		public const int NUL = 0;
		public const int EOS = 1;

		protected virtual void WriteModel(_BinaryWriter bw)
		{
			int c = prob_ep.Count;
			bw.Write(c);
			foreach (var kvp in prob_ep)
			{
				bw.Write7BitEncodedInt(kvp.Key.item);
				bw.Write7BitEncodedInt(kvp.Key.tag);
				bw.Write(kvp.Value);
			}
			new ProbabilityTree { probs = p_tree }.Write(bw);
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// 
		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		Double get_tt_prob(int T1, int T2, int T3)
		{
			ProbabilityTree[] a1, a2;
			return (a1 = p_tree[T1].probs) != null && (a2 = a1[T2].probs) != null ? a2[T3].P : 0.0;
		}
		Double get_bt_prob(int T1, int T2)
		{
			ProbabilityTree[] a1;
			return (a1 = p_tree[T1].probs) != null ? a1[T2].P : 0.0;
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// Decode HMM via Viterbi algorithm.
		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		public pos_prob[][] Decode(int[] inp_words)
		{
#if CLOSED_CLASS_TAGS
			if (closed_class_tags == null)
				setup_closed_classes();
#endif

			int c_tok = inp_words.Length;

			var trellis = new TrellisColumn[c_tok + 1];
			for (int i = 0; i < trellis.Length; i++)
				trellis[i] = new TrellisColumn();

			trellis[0][Bitag.NULEOS] = TrellisData.StartNode;

			for (int i = 0; i < c_tok; i++)
				find_best_extension(inp_words[i], i, trellis);

			return ReadbackPaths(trellis);
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// 
		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		void find_best_extension(int item, int i_col, TrellisColumn[] trellis)
		{
			Bitag default_path = default(Bitag);
			Double sum = 0.0;

			for (int tag = 2; tag < p_tree.Length; tag++)
			{

				TrellisData td;
#if CLOSED_CLASS_TAGS
				if (item == -1 && closed_class_tags[tag])
					td = default(TrellisData);
				else
#endif
				td = maximize_column(new Emission(item, tag), trellis[i_col]);

#if HMM_DEBUGGING
				td.tag = tag;
#endif
				/// ensure that there will be a path forward
				if (td.path.IsNUL)
				{
					if (default_path.IsNUL)
						default_path = trellis[i_col].ArgMax(kvp => kvp.Value.prob).Key;

					/// tiny probability biased by overall probability of trellis tag
					td.prob = pval * p_tree[tag].P;
					td.path = default_path;
				}

				trellis[i_col + 1][new Bitag(td.path.T2, tag)] = td;

				sum += td.prob;
			}

			/// find probabilities
			var dict = trellis[i_col + 1];
			var entr = dict.entries;

			for (int q = 0; q < dict.next_index; q++)
				if (entr[q].hash_code != -1)
					entr[q].value.prob /= sum;
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// 
		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		TrellisData maximize_column(Emission e, TrellisColumn col)
		{
			Double P, TP;
			if (e.item == -1)
				P = p_tree[e.tag].P;
			else if (!prob_ep.TryGetValue(e, out P))
				return default(TrellisData);

			TrellisData td = new TrellisData();
			Debug.Assert(P > 0.0);
			foreach (var td2 in col)
			{
				Debug.Assert(td2.Value.prob != 0.0);

				var bitag = td2.Key;
				if ((TP = get_tt_prob(bitag.T1, bitag.T2, e.tag)) != 0.0)
					TP *= tritag_pref;
				else
				{
					if (bitag.T1 == TrigramHmm.NUL || (TP = get_bt_prob(bitag.T2, e.tag)) == 0.0)
						continue;

					/// backoff from trigram to wildcard-bigram
					TP *= bitag_pref;
					bitag = new Bitag(TrigramHmm.NUL, bitag.T2);
				}
				if ((TP *= td2.Value.prob) > td.prob)
				{
					td.prob = TP;
					td.path = bitag;
				}
			}
			td.prob *= P;
			return td;
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		///
		///////////////////////////////////////////////////////////////////////////////////////////////////////////
		pos_prob[][] ReadbackPaths(TrellisColumn[] trellis)
		{
			int j = trellis.Length - 1;

			pos_prob[][] ret = new pos_prob[j][];

			/// seed the best ending point
			var bt = trellis[j].ArgMax(kvp => kvp.Value.prob).Key;

			/// walk back through trellis and record best path
			while (true)
			{
				j--;

				ret[j] = new pos_prob[1];
				ret[j][0].i_tag = bt.T2;
				ret[j][0].p = 1.0;
				if (j == 0)
					break;

				var btx = trellis[j].ArgMax(kvp => kvp.Value.prob).Key;

				//var btx = trellis[j + 1][bt].path;

				//if (!btf.Equals(btx))
				//	Nop.X();

				bt = btx;
			}
			return ret;
		}
	};

	public struct Emission : IEquatable<Emission>
	{
		public Emission(int item, int tag)
		{
			this.item = item;
			this.tag = tag;
		}
		public int item, tag;
		public override int GetHashCode()
		{
			return (tag << 16) ^ item;
		}
		public override bool Equals(object obj)
		{
			Emission other = (Emission)obj;
			return this.item == other.item && this.tag == other.tag;
		}
		public bool Equals(Emission other)
		{
			return this.item == other.item && this.tag == other.tag;
		}
	};

	public struct Bitag : IEquatable<Bitag>
	{
		public static Bitag NUL = default(Bitag);
		public static Bitag NULEOS = new Bitag(TrigramHmm.NUL, TrigramHmm.EOS);

		public Bitag(int T1, int T2)
			: this()
		{
			this.T1 = T1;
			this.T2 = T2;
		}
		public int T1, T2;
		public bool IsNUL
		{
			get { return T1 == 0 && T2 == 0; }
		}
		public override int GetHashCode()
		{
			return (T1 << 16) ^ T2;
		}
		public override bool Equals(object obj)
		{
			Bitag other = (Bitag)obj;
			return this.T2 == other.T2 && this.T1 == other.T1;
		}
		public bool Equals(Bitag other)
		{
			return this.T2 == other.T2 && this.T1 == other.T1;
		}
#if HMM_DEBUGGING
		public override String ToString()
		{
			var ut = ((StringHmm)TrigramHmm._this).unitags;
			return String.Format("[{0} {1}]", ut[T1], ut[T2]);
		}
#endif
	};

	public struct Tritag : IEquatable<Tritag>
	{
		public Tritag(int ttt, int tt, int t)
		{
			_data = (ttt << 20) | (tt << 10) | t;
		}
		public Tritag(int ttt, Bitag bt)
			: this(ttt, bt.T1, bt.T2)
		{
		}
		int _data;
		public Bitag bt
		{
			get { return new Bitag((int)((uint)_data >> 20), (_data >> 10) & 0x3FF); }
		}
		public int T1 { get { return (int)((uint)_data >> 20); } }
		public int T2 { get { return (_data >> 10) & 0x3FF; } }
		public int T3 { get { return _data & 0x3FF; } }
		public unsafe override int GetHashCode()
		{
			return _data;
		}
		public override bool Equals(object obj)
		{
			return this._data == ((Tritag)obj)._data;
		}
		public bool Equals(Tritag other)
		{
			return this._data == other._data;
		}
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>
	/// 
	/// </summary>
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	[DebuggerDisplay("{ToString(),nq}")]
	public struct ProbabilityTree
	{
		public Double P;
		public ProbabilityTree[] probs;
		public override String ToString()
		{
			return String.Format("{0:N5} {1}",
				P,
				probs == null ? "(none)" : new String(probs.Select(ps => ps.P != 0 ? 'x' : '_').ToArray()));
		}

		public void Write(_BinaryWriter bw)
		{
			bw.Write(P);
			int c = probs == null ? 0 : probs.Length;
			bw.Write(c);
			for (int i = 0; i < c; i++)
				probs[i].Write(bw);
		}

		public void Read(_BinaryReader br)
		{
			this.P = br.ReadDouble();
			int c = br.ReadInt32();
			if (c > 0)
			{
				this.probs = new ProbabilityTree[c];
				for (int i = 0; i < c; i++)
					probs[i].Read(br);
			}
		}
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	///
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	[DebuggerDisplay("{ToString(),nq}")]
	public struct pos_prob : IEquatable<pos_prob>
	{
		public pos_prob(int i_tag, Double p)
		{
			this.i_tag = i_tag;
			this.p = p;
		}
		public int i_tag;
		public Double p;

		public override String ToString()
		{
			return String.Format("{0} {1:N4}", i_tag, p);
		}

		public bool Equals(pos_prob other)
		{
			return this.i_tag == other.i_tag && this.p == other.p;
		}
		public override bool Equals(object obj)
		{
			return this.Equals((PosProb)obj);
		}
		public override int GetHashCode()
		{
			return p.GetHashCode() ^ i_tag;
		}
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	///
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	[DebuggerDisplay("{ToString(),nq}")]
	public struct PosProb : IEquatable<PosProb>
	{
		public PosProb(String pos, Double p)
		{
			this.pos = pos;
			this.p = p;
		}
		public String pos;
		public Double p;

		public override String ToString()
		{
			return String.Format("{0} {1:N4}", pos, p);
		}

		public bool Equals(PosProb other)
		{
			return this.pos == other.pos && this.p == other.p;
		}
		public override bool Equals(object obj)
		{
			return this.Equals((PosProb)obj);
		}
		public override int GetHashCode()
		{
			return p.GetHashCode() ^ pos.GetHashCode();
		}
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>
	/// Wraps the TrigramHmm for use with string-typed inputs and outputs
	/// </summary>
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	public class StringHmm : TrigramHmm
	{
		public StringHmm(
			StringIndex items,
			StringIndex unitags,
			Dictionary<Emission, Double> prob_ep,
			ProbabilityTree[] ps)
			: base(prob_ep, ps)
		{
			this.items = items;
			this.unitags = unitags;
		}
		protected StringHmm(_BinaryReader br)
			: base(br)
		{
			this.items = StringIndex.Read(br);
			this.unitags = StringIndex.Read(br);
		}
		StringHmm(_BinaryReader _ctor_plug, Stream str)
			: this(_ctor_plug = new _BinaryReader(str))
		{
			_ctor_plug.Close();
		}
		public StringHmm(String filename)
			: this(null, new FileStream(filename, FileMode.Open, FileAccess.Read, FileShare.Read))
		{
		}
		public
		readonly StringIndex items, unitags;

		public int GetItemIndex(String item)
		{
			return items[item];
		}
		public String GetTag(int i_tag)
		{
			return unitags[i_tag];
		}
		public PosProb Convert(pos_prob pp)
		{
			return new PosProb(unitags[pp.i_tag], pp.p);
		}

		public PosProb[][] Decode(String[] input)
		{
			/// convert input string items into an array of integer identifiers (unseen values: -1)
			int c = input.Length;
			int[] arr_inp = new int[c];
			for (int i = 0; i < c; i++)
			{
				arr_inp[i] = items[input[i]];
				//if (input[i] == -1)
				//	Console.Write(input[i] + " ");
			}

			var result = base.Decode(arr_inp);

			/// convert answer to the output string items
			var ret = new PosProb[c][];
			for (int i = 0; i < c; i++)
			{
				ret[i] = new PosProb[result[i].Length];
				for (int j = 0; j < ret[i].Length; j++)
				{
#if DEBUG
					if (result[i][j].i_tag <= 1) // NUL, EOS
						throw new Exception();
#endif
					ret[i][j] = Convert(result[i][j]);
				}
			}
			return ret;
		}

		protected override void WriteModel(_BinaryWriter bw)
		{
			base.WriteModel(bw);
			items.Write(bw);
			unitags.Write(bw);
		}
		public void WriteToFile(String filename)
		{
			using (var fs = new FileStream(filename, FileMode.Create, FileAccess.Write, FileShare.None))
			//using (var ds = new DeflateStream(fs, CompressionMode.Compress))
			using (var bw = new _BinaryWriter(fs))
				WriteModel(bw);
		}

		public static StringHmm ReadModel(String filename)
		{
			using (var fs = new FileStream(filename, FileMode.Open, FileAccess.Read, FileShare.Read))
			//using (var ds = new DeflateStream(fs, CompressionMode.Decompress))
			using (var br = new _BinaryReader(fs))
				return new StringHmm(br);
		}

		/// debugging code below here

		public Double Prob(String word, String tag)
		{
			return prob_ep[new Emission(items[word], unitags[tag])];
		}
		public TagProb[] Prob(String word)
		{
			int i_word = items[word];
			return prob_ep.Where(kvp => kvp.Key.item == i_word).Select(kvp => new TagProb { P = kvp.Value, tag = unitags[kvp.Key.tag] }).ToArray();
		}
	};


	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// 
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	[DebuggerDisplay("{ToString(),nq}")]
	public struct TaggedItem
	{
		public TaggedItem(String Item, String Tag)
		{
			this.Item = Item;
			this.Tag = Tag;
		}
		public String Item;
		public String Tag;
		public override String ToString()
		{
			return String.Format("{0} {1}", Item, Tag);
		}
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// 
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	[DebuggerDisplay("{ToString(),nq}")]
	struct tagged_item
	{
		public tagged_item(String s_item, int i_tag)
		{
			this.s_item = s_item;
			this.i_tag = i_tag;
		}
		public String s_item;
		public int i_tag;
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// 
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	public sealed class TrigramHmmTrainer
	{
		public TrigramHmmTrainer()
		{
			this.items = new StringIndex();
			this.unitags = new StringIndex();
			this.τ_unitag = new List<int>();

			if (this.unitags.Add("NUL") != TrigramHmm.NUL)
				throw new Exception();
			τ_unitag.Add(0);

			if (this.unitags.Add("EOS") != TrigramHmm.EOS)
				throw new Exception();
			τ_unitag.Add(0);

			this.τ_itemtag = new Dictionary<Emission, int>();
			this.τ_bitag = new Dictionary<Bitag, int>();
			this.τ_tritag = new Dictionary<Tritag, int>();

			prev_tag = prev_prev_tag = TrigramHmm.EOS;

			c_lines = -1;
			EndSentence();
		}

		const String eos_item = "<eos>";

		/// for the first pass (tallying)
		List<int> τ_unitag;
		Dictionary<Emission, int> τ_itemtag;
		Dictionary<Bitag, int> τ_bitag;
		Dictionary<Tritag, int> τ_tritag;

		int c_obs, c_lines, c_tags, prev_tag, prev_prev_tag;

		/// final products stored in the model
		StringIndex items;
		StringIndex unitags;
		ProbabilityTree[] p_tree;

		public void AddSamples(String file)
		{
			var samples = get_file_sentence_pairs(file);

			tally_samples(samples);
		}

		public void AddSamples(IEnumerable<TaggedItem> rgti)
		{
			tally_samples(rgti);
		}

		public void EndSentence()
		{
			c_lines++;
			tally_sample(eos_item, TrigramHmm.EOS);
			tally_sample(eos_item, TrigramHmm.EOS);
		}

		public StringHmm BuildModel()
		{
			this.c_tags = unitags.Count;

			this.p_tree = new ProbabilityTree[c_tags];

			enhance_tallies();

			Console.WriteLine("sent: {0:#,###}  obs: {1:#,###}  items: {2:#,###}  tags: {3}  tag bigrams: {4}",
					c_lines,
					c_obs,
					items.Count,
					c_tags,
					τ_bitag.Count);
			Console.WriteLine(unitags.StringJoin(" "));
			Console.WriteLine();

			var prob_ep = compute_probabilities();

			return new StringHmm(items, unitags, prob_ep, p_tree);
		}

		IEnumerable<tagged_item> get_file_sentence_pairs(String file)
		{
			using (var sr = new StreamReader(file))
				return get_file_sentence_pairs(sr);
		}

		IEnumerable<tagged_item> get_file_sentence_pairs(StreamReader sr)
		{
			String line;
			while (null != (line = sr.ReadLine()))
			{
				foreach (var ti in get_sentence_pairs(line))
					yield return ti;

				EndSentence();
			}
		}

		IEnumerable<tagged_item> get_sentence_pairs(String line)
		{
			var rgs = line.Split(alib.Character.Charset.tab);
			if ((rgs.Length & 1) != 0)
				throw new Exception(line);

			for (int i = 0; i < rgs.Length; i += 2)
				yield return new tagged_item(rgs[i], unitags.Add(rgs[i + 1]));
		}

		void tally_samples(IEnumerable<TaggedItem> pairs)
		{
			foreach (var ti in pairs)
				tally_sample(ti.Item, unitags.Add(ti.Tag));
		}

		void tally_samples(IEnumerable<tagged_item> pairs)
		{
			foreach (var ti in pairs)
				tally_sample(ti.s_item, ti.i_tag);
		}

		void tally_sample(String s_item, int tag)
		{
			//Console.WriteLine("{0}/{1}", s_item, unitags[tag]);

			c_obs++;
			if (tag == τ_unitag.Count)
				τ_unitag.Add(1);
			else
				τ_unitag[tag]++;

			τ_itemtag.AddIncTally(new Emission(items.Add(s_item), tag));
			τ_bitag.AddIncTally(new Bitag(prev_tag, tag));
			τ_tritag.AddIncTally(new Tritag(prev_prev_tag, prev_tag, tag));

			prev_prev_tag = prev_tag;
			prev_tag = tag;
		}

		void enhance_tallies()
		{
			/// add bitags to the tritag table using the form [NUL TAG TAG]
			int c;
			foreach (var τbt in τ_bitag)
				τ_tritag.Add(new Tritag(TrigramHmm.NUL, τbt.Key), τbt.Value);

			/// add unitags to the bitag table using the form [NUL TAG]
			for (int i = 0; i < c_tags; i++)
			{
				c = τ_unitag[i];
				τ_bitag.Add(new Bitag(TrigramHmm.NUL, i), c);
				p_tree[i].P = c / (Double)c_obs;
			}
		}

		Dictionary<Emission, Double> compute_probabilities()
		{
			int c;
			var prob_ep = new Dictionary<Emission, Double>(τ_itemtag.Count);
			foreach (var τit in τ_itemtag)
				prob_ep.Add(τit.Key, τit.Value / (Double)τ_unitag[τit.Key.tag]);

			int T1;
			foreach (var τbt in τ_bitag)
				if ((T1 = (τbt.Key.T1)) != TrigramHmm.NUL)
					set_bitag_prob(T1, τbt.Key.T2, τbt.Value / (Double)τ_unitag[T1]);

			foreach (var τtt in τ_tritag)
			{
				Tritag tt = τtt.Key;
				if (τ_bitag.TryGetValue(new Bitag(tt.T1, tt.T2), out c))
					set_tritag_prob(tt.T1, tt.T2, tt.T3, τtt.Value / (Double)c);
			}
			return prob_ep;
		}

		void set_bitag_prob(int T1, int T2, Double P)
		{
			ProbabilityTree[] a1;
			if ((a1 = p_tree[T1].probs) == null)
				a1 = p_tree[T1].probs = new ProbabilityTree[c_tags];
			a1[T2].P = P;
		}

		void set_tritag_prob(int T1, int T2, int T3, Double P)
		{
			ProbabilityTree[] a1, a2;
			if ((a1 = p_tree[T1].probs) == null)
				a1 = p_tree[T1].probs = new ProbabilityTree[c_tags];
			if ((a2 = a1[T2].probs) == null)
				a2 = a1[T2].probs = new ProbabilityTree[c_tags];
			a2[T3].P = P;
		}
	};
}
